package org.scohen.juploadr;

import java.lang.instrument.Instrumentation;
import java.lang.reflect.Array;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.IdentityHashMap;
import java.util.Map;

public class SizeOf {

    private static Instrumentation instrumentation;

    public static void premain(String agentArgs, Instrumentation inst) {
        instrumentation = inst;
    }

    public static long sizeOf(Object obj) {
        try {
            return computeObjectSize(obj, new IdentityHashMap<Object, Object>());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static long computeObjectSize(Object obj, Map<Object, Object> visited) throws Exception {
        if (obj == null || visited.containsKey(obj))
            return 0;
        long size = instrumentation.getObjectSize(obj);
        visited.put(obj, visited);
        Class c = obj.getClass();
        if (c.isArray()) {
            if (!c.getComponentType().isPrimitive()) {
                int len = Array.getLength(obj);
                for (int i = 0; i < len; ++i) {
                    size += computeObjectSize(Array.get(obj, i), visited);
                }
            }
        } else {
            for (; c != null; c = c.getSuperclass()) {
                for (Field f : c.getFields()) {
                    if (f.getType().isPrimitive() || Modifier.isStatic(f.getModifiers()))
                        continue;
                    f.setAccessible(true);
                    size += computeObjectSize(f.get(obj), visited);
                }
            }
        }
        return size;
    }
}
